#include "avl_tree.h"
#include <stdio.h>
#include <stdlib.h>

/**
 * @brief new_avl_tree_node
 * @param data
 * @return avl tree node
 */
tree_node_t *new_avl_tree_node(int data){
    tree_node_t *node = (tree_node_t *)malloc(sizeof(tree_node_t));
    if(node == NULL){
        return NULL;
    }
    node->data = data;
    node->left = NULL;
    node->right = NULL;
    node->height = 1;

    return node;
}

/**
 * @brief free_avl_tree_node
 * @param node
 */
void free_avl_tree_node(tree_node_t *node){
    if(node == NULL){
        return;
    }
    free(node);
}

/**
 * @brief get_balance_factor
 * @param node
 * @return
 */
int get_balance_factor(tree_node_t *node){
    if(node == NULL){
        return 0;
    }
    int left_factor = 0;
    int right_factor = 0;

    if(node->left != NULL){
        left_factor = node->left->height;
    }
    if(node->right != NULL){
        right_factor = node->right->height;
    }

    return left_factor - right_factor;
}

/**
 * @brief avl_tree_height_update
 * @param root
 */
void avl_tree_height_update(tree_node_t *root){
    if(root == NULL){
        return;
    }

    if(root->left != NULL && root->right != NULL){
        //结点有左右子树
        //左右子树存在则获取最大子树的高度
        root->height = root->left->height > root->right->height ? root->left->height : root->right->height;
        root->height += 1;
    }else if(root->left != NULL){
        //只存在左子树
        root->height = root->left->height;
        root->height += 1;
    }else if(root->right != NULL){
        //只存在右子树
        root->height = root->right->height;
        root->height += 1;
    }

    //printf("current node height: %d\n", root->height);
}

int get_max_height(tree_node_t *node){
    if(node == NULL){
        return 0;
    }
    if(node->left != NULL && node->right != NULL){
        return node->left->height > node->right->height ? node->left->height : node->right->height;
    }else if(node->left != NULL){
        return node->left->height;
    }else if(node->right != NULL){
        return node->right->height;
    }
    return node->height;
}

/**
 * @brief avl_tree_right_rotate
 * @param x
 * @return
 */
tree_node_t *avl_tree_right_rotate(tree_node_t *x){
    tree_node_t *y = x->left;
    tree_node_t *t2 = y->right;

    y->right = x;
    x->left = t2;

    //右旋后必须先求出x的高度再求出y的高度
    int max_height = get_max_height(x);
    x->height = max_height + 1;
    max_height = get_max_height(y);
    y->height = max_height + 1;

    return y;
}

/**
 * @brief avl_tree_left_rotate
 * @param x
 * @return
 */
tree_node_t *avl_tree_left_rotate(tree_node_t *x){
    tree_node_t *y = x->right;
    tree_node_t *t2 = y->left;

    y->left = x;
    x->right = t2;

    //右旋后必须先求出x的高度再求出y的高度
    int max_height = get_max_height(x);
    x->height = max_height + 1;
    max_height = get_max_height(y);
    y->height = max_height + 1;

    return y;
}

/**
 * @brief avl_tree_balance_adjust
 * @param root_node
 * @return
 */
tree_node_t *avl_tree_balance_adjust(tree_node_t *root_node){
    if(root_node == NULL){
        return NULL;
    }

    avl_tree_height_update(root_node);  //更新结点高度
    int balance_factor = get_balance_factor(root_node);
    if(balance_factor > 1 && get_balance_factor(root_node->left) > 0){
        //左子树-右子树高度差大于1且root_node左孩子的平衡因子大于0，则是LL场景，需要右旋
        printf("need ll\n");
        return avl_tree_right_rotate(root_node);
    } else if(balance_factor > 1 && get_balance_factor(root_node->left) < 0){
        //左子树-右子树高度差大于1且root_node左孩子的平衡影子小于0，则是LR场景，需要先左旋再右旋
        printf("need lr\n");
        root_node->left = avl_tree_left_rotate(root_node->left);
        return avl_tree_right_rotate(root_node);
    }else if(balance_factor < -1 && get_balance_factor(root_node->right) < 0){
        //左子树-右子树高度差小于-1且root_node右孩子平衡因子大于0，则是RR场景，需要左旋
        printf("need rr\n");
        return avl_tree_left_rotate(root_node);
    }else if(balance_factor < -1 && get_balance_factor(root_node->right) > 0){
        //左子树-右子树高度差小于-1且root_node右孩子平衡因子小于0，则是RL场景，需要先右旋再左旋
        printf("need rl\n");
        root_node->right = avl_tree_right_rotate(root_node->right);
        return avl_tree_left_rotate(root_node);
    }

    return root_node;
}

/**
 * @brief sub_tree_add_node
 * @param root_node
 * @param new_node
 * @return
 */
tree_node_t *sub_tree_add_node(tree_node_t *root_node, int data){
    if(root_node == NULL){
        return new_avl_tree_node(data);
    }
    //找到新结点适合插入的位置
    if(data < root_node->data){
        root_node->left = sub_tree_add_node(root_node->left, data);
    }else{
        root_node->right = sub_tree_add_node(root_node->right, data);
    }
    //printf("insert node: %d\n", data);
    //printf("current node: %d, height = %d\n", root_node->data, root_node->height);

    return avl_tree_balance_adjust(root_node);  //调节avl树
}

/**
 * @brief avl_tree_insert_node
 * @param tree root
 * @param data
 */
void avl_tree_insert_node(tree_node_t **root, int data){
    if(root == NULL){
        return;
    }
    tree_node_t *tree_root = *root;
    tree_root = sub_tree_add_node(tree_root, data);  //将新结点放入某个合适位置，后面要调整平衡性
    *root = tree_root;
}

/**
 * @brief sub_tree_delete_node
 * @param root_node
 * @param data
 * @return
 */
tree_node_t *sub_tree_delete_node(tree_node_t *root_node, int data){
    if(root_node == NULL){
        return NULL;
    }

    tree_node_t *node = NULL;
    if(data < root_node->data){
        root_node->left = sub_tree_delete_node(root_node->left, data);
    }else if(data > root_node->data){
        root_node->right = sub_tree_delete_node(root_node->right, data);
    }else{
        //找到要删除的结点
        if(root_node->left == NULL && root_node->right != NULL){
            //左子树为空，则用右子树结点代替被删除结点
            tree_node_t *right = root_node->right;
            root_node->right = NULL;
            free_avl_tree_node(root_node);
            return right;
        }else if(root_node->right == NULL && root_node->left != NULL){
            //右子树为空，则用左子树结点代替被删除结点
            tree_node_t *left = root_node->left;
            root_node->left = NULL;
            free_avl_tree_node(root_node);
            return left;
        }else if(root_node->right == NULL && root_node->left == NULL){
            free_avl_tree_node(root_node);
            return NULL;
        }
        else{
            //左右子树不为空的情况
            tree_node_t *temp_node = NULL;
            if(get_balance_factor(root_node) >= 0){
                //左子树高度大于等于右子树高度，采用左子树结点代替原结点
                node = root_node->left;
                temp_node = node->right;
                node->right = root_node->right;
                free_avl_tree_node(root_node);
                return sub_tree_add_node(node, temp_node);
            }else{
                //右子树高度大于左子树，采用右子树结点代替原结点
                node = root_node->right;
                temp_node = node->left;
                node->left = root_node->left;
                free_avl_tree_node(root_node);
                return sub_tree_add_node(node, temp_node);
            }
        }
    }
    return root_node;
}

/**
 * @brief avl_tree_balance_adjust_all
 * @param node
 * @return
 */
tree_node_t *avl_tree_balance_adjust_all(tree_node_t *node){
    if(node == NULL){
        return NULL;
    }
    avl_tree_balance_adjust_all(node->left);
    avl_tree_balance_adjust_all(node->right);
    return avl_tree_balance_adjust(node);
}

/**
 * @brief avl_tree_delete_node
 * @param root
 * @param data
 */
void avl_tree_delete_node(tree_node_t **root, int data){
    if(root == NULL){
        return;
    }
    if(*root == NULL){
        return;
    }

    *root = sub_tree_delete_node(*root, data);
    tree_node_t *root_node = *root;
    *root = avl_tree_balance_adjust_all(root_node);
}

/**
 * @brief avl_tree_print
 * @param root
 */
void avl_tree_print(tree_node_t *root){
    if(root == NULL){
        return;
    }
    printf("%d，", root->data);
    avl_tree_print(root->left);
    avl_tree_print(root->right);
}
